#ifndef NE_STL_INCLUDENE_STL_SHARED_PTR_H
#define NE_STL_INCLUDENE_STL_SHARED_PTR_H

#include "ne_stl_build.h"
#include "ne_stl_memory.h"

namespace nstd {

namespace {

typedef struct _cblock {
    _cblock() {}
    nstd::atomic_int weak_count{0};
    nstd::atomic_int shared_count{1};

} cblock;
void cblock_release(cblock*& ptr) {
    memory::delete_obj(ptr);
}
template <class TObj>
void object_release(TObj*& ptr) {
    memory::delete_obj(ptr);
}
}  // namespace

template <class T>
class _weak_ptr;

template <class T>
class _shared_ptr {
    template <class TObj>
    friend class _weak_ptr;
    template <class TObj>
    friend class _shared_ptr;

public:
    using element_type = std::remove_extent_t<T>;
    _shared_ptr()
        : data_(nullptr)
        , release_(object_release<T>)
        , release_block_(cblock_release) {}

    explicit _shared_ptr(T* data)
        : data_(data)
        , release_(object_release<T>)
        , release_block_(cblock_release) {
        if (data_) {
            cblock_ = memory::new_obj<cblock>();
        }
    }
    _shared_ptr(std::nullptr_t)
        : data_(nullptr)
        , release_(object_release<T>)
        , release_block_(cblock_release) {}
    _shared_ptr(const _shared_ptr& sp)
        : data_(sp.data_)
        , cblock_(sp.cblock_)
        , release_(sp.release_)
        , release_block_(sp.release_block_) {
        if (cblock_) {
            cblock_->shared_count.operator++();
        }
    }
    template <class D>
    _shared_ptr(const _shared_ptr<D>& sp)
        : data_(dynamic_cast<T*>(sp.data_))
        , cblock_(sp.cblock_)
        , release_(reinterpret_cast<void (*)(T*&)>(sp.release_))
        , release_block_(sp.release_block_) {
        if (cblock_) {
            cblock_->shared_count.operator++();
        }
    }
    _shared_ptr(const _weak_ptr<T>& wp)
        : data_(wp.data_)
        , cblock_(wp.cblock_)
        , release_(wp.release_)
        , release_block_(cblock_release) {
        if (cblock_) {
            cblock_->shared_count.operator++();
        }
    }
    ~_shared_ptr() { dec_shared_count(); }
    _shared_ptr& operator=(const _shared_ptr& sp) {
        if (this == &sp)
            return *this;
        data_ = sp.data_;
        cblock_ = sp.cblock_;
        release_ = sp.release_;
        release_block_ = sp.release_block_;
        if (cblock_) {
            cblock_->shared_count.operator++();
        }
        return *this;
    }
    _shared_ptr& operator=(std::nullptr_t) {
        reset();
        return *this;
    }
    inline operator bool() const { return data_ != nullptr; }
    inline bool operator==(std::nullptr_t) const { return !this->operator bool(); }
    inline bool operator!=(std::nullptr_t) const { return this->operator bool(); }
    T* get() const { return data_; }

    T* operator->() const { return data_; }

    T& operator*() const {
        assert(*this);
        return *data_;
    }

    void reset() {
        dec_shared_count();
        data_ = nullptr;
        cblock_ = nullptr;
    }

    void reset(T* data) {
        dec_shared_count();
        data_ = data;
        cblock_ = memory::new_obj<cblock>();
    }

private:
    inline void dec_shared_count() {
        if (!cblock_) {
            return;
        }
        cblock_->shared_count.operator--();
        if (cblock_->shared_count <= 0) {
            if (release_ != nullptr)
                release_(data_);
            if (cblock_->weak_count <= 0 && release_block_ != nullptr) {
                release_block_(cblock_);
            }
        }
    }

private:
    T* data_{nullptr};
    nstd::cblock* cblock_{nullptr};
    void (*release_)(T*& pdata){nullptr};
    void (*release_block_)(nstd::cblock*& pdata){nullptr};
};
template <class T>
class _weak_ptr {
public:
    _weak_ptr()
        : data_(nullptr)
        , cblock_(nullptr)
        , release_(nullptr)
        , release_block_(nullptr) {}
    _weak_ptr(const _weak_ptr& wp)
        : data_(wp.data_)
        , cblock_(wp.cblock_)
        , release_(wp.release_)
        , release_block_(wp.release_block_) {
        if (cblock_) {
            cblock_->weak_count.operator++();
        }
    }
    _weak_ptr(const _shared_ptr<T>& sp)
        : data_(sp.data_)
        , cblock_(sp.cblock_)
        , release_(sp.release_)
        , release_block_(sp.release_block_) {
        if (cblock_) {
            cblock_->weak_count.operator++();
        }
    }
    ~_weak_ptr() { dec_weak_count(); }
    _weak_ptr& operator=(const _weak_ptr& wp) {
        _weak_ptr tmp(wp);
        swap(tmp);
        return *this;
    }
    inline _shared_ptr<T> lock() const { return expired() ? _shared_ptr<T>(nullptr) : _shared_ptr<T>(*this); }
    inline bool expired() const { return cblock_ == nullptr || cblock_->shared_count <= 0; }
    inline void reset() {
        dec_weak_count();
        data_ = nullptr;
        cblock_ = nullptr;
        release_block_ = nullptr;
    }

private:
    inline void dec_weak_count() {
        if (cblock_ == nullptr) {
            return;
        }
        cblock_->weak_count.operator--();
        if (cblock_->shared_count <= 0 && cblock_->weak_count <= 0 && release_block_ != nullptr) {
            release_block_(cblock_);
            cblock_ = nullptr;
        }
    }

private:
    T* data_;
    nstd::cblock* cblock_;
    void (*release_)(T*& pdata);
    void (*release_block_)(nstd::cblock*& pdata);
    template <class Y>
    friend class _shared_ptr;
};
template <class T>
using shared_ptr = nstd::_shared_ptr<T>;
template <class T, class... TArgs>
nstd::shared_ptr<T> make_shared(TArgs&&... args) {
    return nstd::shared_ptr<T>(memory::new_obj<T>(args...));
}
template <class T1, class T2>
shared_ptr<T1> dynamic_pointer_cast(const shared_ptr<T2>& other) noexcept {
    // dynamic_cast for shared_ptr that properly respects the reference count control block
    const auto ptr = dynamic_cast<typename shared_ptr<T1>::element_type*>(other.get());
    if (ptr) {
        return shared_ptr<T1>(other);
    }
    return (shared_ptr<T1>());
}
template <class T>
using weak_ptr = nstd::_weak_ptr<T>;
}  // namespace nstd

#endif  // !NE_STL_INCLUDENE_STL_SHARED_PTR_H